package com.andnnl.ratelimit;


import cn.hutool.core.thread.ThreadUtil;
import com.google.common.util.concurrent.RateLimiter;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 做接口访问限制，访问一次获取一批数据，如果每个用户访问有数据，每秒最多访问10次，如果没数据，直接返回空，1秒最多访问1次，接口支持多用户同时访问的，每个用户需要自己的RateLimiter
 * Created by chenss on 2024/3/18.
 */
public class SdkAccessLimiter {
    private Map<String, RateLimiter> rateLimitersWithData = new ConcurrentHashMap<>();//有数据的用户限制，每秒10次
    private Map<String, RateLimiter> rateLimitersWithoutData = new ConcurrentHashMap<>();//没有数据的用户限制，每秒2次
    private Map<String, Boolean> lastHasDataMap = new HashMap<>();//记录上次有没有数据
    private int permitsWithData = 10;
    private int permitsWithoutData = 2;

    public SdkAccessLimiter() {
    }

    public SdkAccessLimiter(int permitsWithData, int permitsWithoutData) {
        this.permitsWithData = permitsWithData;
        this.permitsWithoutData = permitsWithoutData;
    }

    /**
     * 尝试本次是否能访问，立即返回的
     *
     * @param userId
     * @return
     */
    public boolean tryAccess(String userId) {
        boolean lastHasData = lastHasDataMap.getOrDefault(userId, false);
        RateLimiter limiter;
        if (lastHasData) {
            limiter = rateLimitersWithData.computeIfAbsent(userId, k -> RateLimiter.create(permitsWithData));
        } else {
            limiter = rateLimitersWithoutData.computeIfAbsent(userId, k -> RateLimiter.create(permitsWithoutData));
        }
        return limiter.tryAcquire();
    }

    /**
     * 设置上次有没有数据
     *
     * @param userId
     * @param hasData
     */
    public void putHashData(String userId, boolean hasData) {
        lastHasDataMap.put(userId, hasData);
    }


    public static void main(String[] args) {
        SdkAccessLimiter limiter = new SdkAccessLimiter();
        // 模拟用户访问接口
        for (int i = 0; i < 20; i++) {
            String userId = "user" + i % 2;
            if (limiter.tryAccess(userId)) {
                boolean hasData = i % 2 == 0;
                System.out.println(System.currentTimeMillis() + " " + userId + " accessed the API " + hasData);
                limiter.putHashData(userId, hasData);
            } else {
                System.out.println(System.currentTimeMillis() + " " + userId + " failed to access the API" );
            }
            ThreadUtil.sleep(200);
        }
    }
}
